# -*- coding: utf-8 -*-
"""
TCP 端口映射工具
提供双向数据转发功能，用于 FRP 客户端和服务器之间的连接转发
"""

import sys
import socket
import logging
import threading

# 接收数据缓存大小
PKT_BUFF_SIZE = 10240

# 获取主程序的日志记录器
logger = logging.getLogger("FRP")
formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s", "%H:%M:%S")
# 只输出到控制台，不记录到文件
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)

# 只显示 INFO 及以上级别的日志
logger.setLevel(logging.INFO)


def tcp_mapping_worker(
    conn_receiver: socket.socket, conn_sender: socket.socket
) -> None:
    """
    单向流数据传递工作线程

    参数:
        conn_receiver: 数据接收端连接
        conn_sender: 数据发送端连接
    """
    try:
        while True:
            try:
                data = conn_receiver.recv(PKT_BUFF_SIZE)
            except Exception:
                break

            if not data:
                break

            try:
                conn_sender.sendall(data)
            except Exception:
                break

    finally:
        try:
            conn_receiver.close()
        except:
            pass
        try:
            conn_sender.close()
        except:
            pass


def join(conn_a: socket.socket, conn_b: socket.socket) -> None:
    """
    建立双向数据转发

    参数:
        conn_a: 第一个连接
        conn_b: 第二个连接
    """
    threading.Thread(
        target=tcp_mapping_worker, args=(conn_a, conn_b), daemon=True
    ).start()
    threading.Thread(
        target=tcp_mapping_worker, args=(conn_b, conn_a), daemon=True
    ).start()


if __name__ == "__main__":
    # 测试代码
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(("0.0.0.0", 8080))
    sock.listen(100)
    conn1, addr1 = sock.accept()

    sock2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock2.bind(("0.0.0.0", 7000))
    sock2.listen(100)

    # 类似frpc
    server_conn = socket.create_connection(("localhost", 7000))

    conn2, addr2 = sock2.accept()
    join(conn1, conn2)

    target_conn = socket.create_connection(("192.168.1.101", 3389))
    join(server_conn, target_conn)
